# -*- coding: UTF-8 -*-
"""
@Project ：ai-demo
@File    ：asr_llm_tts_pipeline_full.py
@IDE     ：PyCharm
@Author  ：zhangyiheng
@Date    ：2025/4/1 12:04
@Describe：ASR + VAD + LLM + TTS
"""
import os
import re
import io
import queue
import time
import threading
import requests
import numpy as np
from datetime import datetime
from pathlib import Path
from typing import Optional, Union

import webrtcvad
from scipy.signal import resample

try:
    import sounddevice as sd
    import soundfile as sf
except ImportError:
    raise ImportError("需要安装 sounddevice 和 soundfile 库：pip install sounddevice soundfile")

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.outputs import OutputKeys
from funasr import AutoModel
import logging
#from commcon.constant import MODEL_ROOT
log = logging.getLogger('modelscope')

# 配置参数
MODEL_ROOT = "/Users/zhangzhan/Desktop/zhangzhan/modelscopes/repository-c"
OLLAMA_HOST = "http://localhost:11434"
OLLAMA_MODEL = "qwen2.5:1.5b"
SAMPLE_RATE = 16000
CHUNK_SIZE = 2048
SILENCE_TIMEOUT = 2  # 静音超时(秒)
MIN_RECORDING = 0.8  # 最短有效录音(秒)


class AudioRecorder:
    """带VAD的智能录音器"""

    def __init__(self):
        self.sample_rate = SAMPLE_RATE
        self.chunk_size = CHUNK_SIZE
        self.silence_timeout = SILENCE_TIMEOUT
        self.min_recording = MIN_RECORDING
        self.recording = False
        self.last_voice_time = 0
        self.audio_buffer = []
        self.current_audio = None
        self.energy_threshold = 0.02
        self.noise_floor = 0.01
        self.running = True
        self.stream = None

        # 新增WebRTC-VAD初始化
        self.vad = webrtcvad.Vad(2)  # 激进模式（0~3）
        self.vad_sample_rate = 16000  # WebRTC-VAD固定要求16kHz

        self.playback_active = False  # 播放状态标记

    def _is_human_voice(self, audio_data):
        """WebRTC-VAD检测是否为人声"""
        # 1. 确保音频长度足够（至少30ms）
        required_length = int(0.03 * self.vad_sample_rate)  # 480 samples @16kHz
        if len(audio_data) < required_length:
            return False  # 或补零: audio_data = np.pad(audio_data, (0, required_length - len(audio_data)))

        # 2. 重采样至16kHz（更安全的实现）
        if self.sample_rate != self.vad_sample_rate:
            from scipy.signal import resample_poly  # 更稳定的重采样方法
            audio_data = resample_poly(audio_data, self.vad_sample_rate, self.sample_rate)

        # 3. 转换到int16（带溢出保护）
        audio_data = np.clip(audio_data, -1.0, 1.0)  # 限制在[-1, 1]范围内
        audio_int16 = (audio_data * 32767).astype(np.int16)

        # 4. 确保字节长度正确（30ms的帧）
        frame_bytes = audio_int16[:required_length].tobytes()  # 严格截取480采样点

        try:
            return self.vad.is_speech(frame_bytes, self.vad_sample_rate)
        except Exception as e:
            log.error(f"VAD处理异常: {str(e)}")
            return False  # 出错时保守返回False

    def _update_threshold(self, audio_data):
        """动态调整VAD阈值"""
        current_energy = np.mean(np.abs(audio_data))
        if current_energy < self.noise_floor * 1.5:
            self.energy_threshold = self.noise_floor * 2
        else:
            self.energy_threshold = max(current_energy * 0.7, self.noise_floor * 2)

    def _audio_callback(self, indata, frames, _time, status):
        """音频处理回调（集成VAD）"""

        if self.playback_active:
            return

        audio_data = np.frombuffer(indata, dtype=np.float32)
        self._update_threshold(audio_data)

        # 语音活动检测（能量+WebRTC-VAD双验证）
        current_energy = np.mean(np.abs(audio_data))
        is_speaking = (current_energy > self.energy_threshold) and \
                     self._is_human_voice(audio_data)  # 新增VAD检测

        if is_speaking:
            sd.stop()  # 立即停止所有播放
            time.sleep(0.08)  # 等待音频设备复位
            self.last_voice_time = _time.currentTime
            if not self.recording:
                self._start_recording()

        if self.recording:
            self.audio_buffer.append(audio_data.copy())
            silent_duration = _time.currentTime - self.last_voice_time
            if silent_duration > self.silence_timeout:
                self._stop_recording()

    def _start_recording(self):
        """开始录音"""
        self.recording = True
        self.audio_buffer = []
        log.warning("🔊 检测到语音，开始录音...")

    def _stop_recording(self):
        """停止录音"""
        self.recording = False
        duration = len(self.audio_buffer) * self.chunk_size / self.sample_rate

        if duration >= self.min_recording:
            log.warning(f"🛑 录音结束 (时长: {duration:.1f}s)")
            self.current_audio = self._save_recording()
        else:
            log.warning("❌ 忽略过短录音")
            self.audio_buffer = []

    def _save_recording(self):
        """保存录音文件"""
        audio_data = np.concatenate(self.audio_buffer)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        temp_path = f"temp_recording_{timestamp}.wav"
        sf.write(temp_path, audio_data, self.sample_rate)
        return temp_path

    def start(self):
        """启动录音监听"""

        def _run():
            with sd.InputStream(
                    callback=self._audio_callback,
                    channels=1,
                    samplerate=self.sample_rate,
                    blocksize=self.chunk_size,
                    dtype='float32'
            ):
                while self.running:
                    time.sleep(0.1)

        self.thread = threading.Thread(target=_run, daemon=True)
        self.thread.start()
        log.info("🎤 麦克风监听已就绪...")

    def stop(self):
        """停止录音"""
        self.running = False
        if self.thread:
            self.thread.join()


class OrderedAudioPlayer:
    """保序音频播放器（支持中断）"""

    def __init__(self):
        self.play_queue = queue.PriorityQueue()
        #self.current_seq = 0
        self.play_thread = threading.Thread(target=self._play_worker, daemon=True)
        self.lock = threading.Lock()
        self.is_running = True
        self.condition = threading.Condition()
        self._active_request_id = 0  # 当前有效请求ID
        self.play_thread.start()

        self.expected_seq = 0  # 当前期待播放的序列号
        self.current_req_id = 0  # 当前处理的请求ID

    def _play_worker(self):
        """播放工作线程"""
        while self.is_running or not self.play_queue.empty():
            with self.condition:
                try:
                    priority, (seq, audio_data, req_id) = self.play_queue.get_nowait()

                    # 有效性检查：只处理当前最新请求
                    if req_id != self._active_request_id:
                        continue

                    # 新请求重置序列号
                    if req_id != self.current_req_id:
                        self.current_req_id = req_id
                        self.expected_seq = 0

                    if seq == self.expected_seq:
                        self._play_audio(audio_data)
                        self.expected_seq += 1
                    else:
                        # 重新入队等待正确顺序
                        self.play_queue.put((priority, (seq, audio_data, req_id)))
                        self.condition.wait(0.1)
                except queue.Empty:
                    self.condition.wait(0.5)

    def _play_audio(self, audio_data: bytes):
        """实际播放音频（增加状态通知）"""

        with io.BytesIO(audio_data) as f:
            data, sr = sf.read(f)
            try:
                sd.play(data, sr)
                sd.wait()
            except Exception as e:
                log.error(f"播放异常: {str(e)}")
                raise

    def add_audio(self, audio_data: bytes, seq: int, request_id: int):
        """添加音频到队列"""
        with self.lock:
            self.play_queue.put((seq, (seq, audio_data, request_id)))
            with self.condition:
                self.condition.notify()

    def clear_queue(self, new_request_id: int):
        """清空队列并更新请求ID"""
        with self.lock:
            self._active_request_id = new_request_id
            while not self.play_queue.empty():
                try:
                    self.play_queue.get_nowait()
                except queue.Empty:
                    break
            sd.stop()  # 立即停止当前播放
            with self.condition:
                self.condition.notify_all()
            log.warning("🔄 播放队列已重置")

    def stop(self):
        """停止播放器"""
        self.is_running = False
        with self.condition:
            self.condition.notify_all()
        self.play_thread.join()


class LLMClient:
    """Ollama大模型客户端"""

    def __init__(self):
        self.session = requests.Session()
        self.session.timeout = 30

    def generate(self, prompt: str) -> Optional[str]:
        """生成回复"""
        try:
            response = self.session.post(
                f"{OLLAMA_HOST}/api/generate",
                json={
                    "model": OLLAMA_MODEL,
                    "prompt": prompt,
                    "stream": False,
                    "options": {"temperature": 0.7}
                }
            )
            response.raise_for_status()
            return self._clean_text(response.json()["response"])
        except Exception as e:
            log.error(f"LLM请求失败: {str(e)}")
            return None

    def _clean_text(self, text: str) -> str:
        """清洗文本输出"""
        text = re.sub(r'\s+', ' ', text)  # 合并空格
        text = re.sub(r'[\n\t]', ' ', text)  # 去除换行符
        return text.strip()


class SpeechProcessor:
    """语音处理核心（支持中断）"""

    def __init__(self, max_workers=8):
        # 初始化组件
        self.asr_model = self._init_asr_model()
        self.tts_pipeline = self._init_tts_pipeline()
        self.llm = LLMClient()
        self.audio_player = OrderedAudioPlayer()
        self.recorder = AudioRecorder()

        # 线程控制
        self.semaphore = threading.Semaphore(max_workers)
        self.request_counter = 0
        self.request_lock = threading.Lock()
        self.is_running = True

    def _init_asr_model(self):
        """初始化ASR模型"""
        return AutoModel(
            model=f"{MODEL_ROOT}/iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
            model_revision="v2.0.5",
            vad_model=f"{MODEL_ROOT}/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
            vad_model_revision="v2.0.4",
            punc_model=f"{MODEL_ROOT}/iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch",
            punc_model_revision="v2.0.4",
            disable_update=True,
            cache_dir=MODEL_ROOT,
        )

    def _init_tts_pipeline(self):
        """初始化TTS管道"""
        return pipeline(
            task=Tasks.text_to_speech,
            model=f"{MODEL_ROOT}/iic/speech_sambert-hifigan_tts_zh-cn_16k",
            model_revision="v1.0.2"
        )

    def start(self):
        """启动系统"""
        self.recorder.start()
        try:
            while self.is_running:
                if self.recorder.current_audio:
                    audio_path = self.recorder.current_audio
                    self.recorder.current_audio = None

                    # 启动新请求处理
                    threading.Thread(
                        target=self._process_pipeline,
                        args=(audio_path,),
                        daemon=True
                    ).start()

                time.sleep(0.1)
        except KeyboardInterrupt:
            log.warning("正在关闭系统...")
        finally:
            self.stop()

    def _process_pipeline(self, audio_path: str):
        """完整处理流水线"""
        try:
            # 1. ASR识别（先不中断前序播放）
            asr_text = self._speech_to_text(audio_path)

            # 有效性检查
            if not asr_text or self._is_ambient_noise(asr_text):  # 新增环境噪音检测
                log.warning("🛂 忽略无效语音输入")
                return

            log.info(f"\n🗣 有效输入: {asr_text}")

            # 生成新请求ID
            with self.request_lock:
                self.request_counter += 1
                current_request_id = self.request_counter

            # 2. 确认有效输入后才中断前序处理
            self.audio_player.clear_queue(current_request_id)

            # LLM处理
            llm_response = self.llm.generate(asr_text)
            if not llm_response:
                return

            log.info(f"🧠 生成回复: {llm_response}")

            # TTS输出
            self._text_to_speech(llm_response, current_request_id)

        finally:
            try:
                os.remove(audio_path)
            except:
                pass

    def _is_ambient_noise(self, text: str) -> bool:
        """环境噪音检测"""
        # 规则1：排除纯标点/语气词
        if re.fullmatch(r'[\s,.!?嗯啊哦]+', text):
            return True

        # 规则2：排除过短无意义词
        if len(text.strip()) < 3 and not any(c.isalnum() for c in text):
            return True

        # 规则3：常见误识别模式（如"嗯"、"啊"等）
        noise_patterns = [r'^嗯+$', r'^啊+$', r'^呃+$']
        if any(re.match(p, text) for p in noise_patterns):
            return True

        return False

    def _speech_to_text(self, audio_path: str) -> Optional[str]:
        """语音识别"""
        try:
            result = self.asr_model.generate(input=audio_path)
            return self._extract_asr_text(result)
        except Exception as e:
            log.error(f"ASR识别失败: {str(e)}")
            return None

    def _extract_asr_text(self, result) -> str:
        """提取ASR文本"""
        if isinstance(result, list):
            return result[0].get("text", "").strip()
        elif isinstance(result, dict):
            return result.get("text", "").strip()
        return str(result).strip()

    def _text_to_speech(self, text: str, request_id: int):
        """文本转语音（每个请求独立序列号）"""

        # 播放开始标记
        self.recorder.playback_active = True

        try:
            sentences = self._split_sentences(text)
            for sentence_seq, sentence in enumerate(sentences):  # 请求内独立序列号
                threading.Thread(
                    target=self._tts_worker,
                    args=(sentence, sentence_seq, request_id),  # 使用局部序列号
                    daemon=True
                ).start()
        finally:
            threading.Thread(
                target=self._enable_recording_delayed,
                args=(request_id,),
                daemon=True
            ).start()

    def _enable_recording_delayed(self, request_id: int):
        """延迟恢复录音检测"""
        time.sleep(0.5)  # 等待播放完全结束
        with self.request_lock:
            if self.request_counter == request_id:  # 仍是当前有效请求
                self.recorder.playback_active = False

    def _tts_worker(self, text: str, seq: int, req_id: int):
        """TTS工作线程"""
        try:
            with self.semaphore:
                # 检查请求有效性
                with self.request_lock:
                    if self.request_counter != req_id:
                        log.warning(f"[请求#{req_id}] 已过期，丢弃句子")
                        return

                output = self.tts_pipeline(input=text, voice="zhiyan_emo")

                # 再次检查有效性
                with self.request_lock:
                    if self.request_counter != req_id:
                        log.warning(f"[请求#{req_id}] 已过期，丢弃生成的音频")
                        return

                self.audio_player.add_audio(output[OutputKeys.OUTPUT_WAV], seq, req_id)
        except Exception as e:
            log.error(f"[句子#{seq}] TTS失败: {str(e)}")

    def _split_sentences(self, text: str) -> list:
        """智能分段"""
        sentences = re.split(r'(?<=[。！？.!?；])', text)
        merged = []
        buffer = ""

        for s in [s.strip() for s in sentences if s.strip()]:
            if len(buffer + s) < 15:  # 合并短句
                buffer += s
            else:
                if buffer:
                    merged.append(buffer)
                buffer = s

        if buffer:
            merged.append(buffer)

        return merged



    def stop(self):
        """停止系统"""
        self.is_running = False
        self.recorder.stop()
        self.audio_player.stop()


def check_audio_setup():
    """音频设备检查"""
    log.warning("=== 音频设备检查 ===")
    try:
        log.warning("可用输入设备：")
        log.warning(sd.query_devices(kind='input'))

        # 测试录音
        log.warning("正在测试麦克风...")
        test_rec = sd.rec(int(0.5 * SAMPLE_RATE), samplerate=SAMPLE_RATE, channels=1)
        sd.wait()
        log.info("✅ 麦克风测试成功")

        # 测试播放
        log.warning("正在测试扬声器...")
        sd.play(test_rec, SAMPLE_RATE)
        sd.wait()
        log.info("✅ 扬声器测试成功")
    except Exception as e:
        log.error(f"❌ 音频设备异常: {str(e)}")
        raise


if __name__ == "__main__":
    check_audio_setup()
    log.warning("=== 启动语音交互系统 ===")
    processor = SpeechProcessor(max_workers=3)

    try:
        processor.start()
    except Exception as e:
        log.error(f"系统异常: {str(e)}")
    finally:
        processor.stop()